from pyspark import SparkContext
from pyspark.sql import SparkSession
import json
import numpy as np
import pandas as pd
sc = SparkContext("local", "duty 1")
columns_json_str = '{"date":"date","stock_symbol":"stock_id","stock_volume":"volume"}'
columns_dict = json.loads(columns_json_str)
sc.setLogLevel('WARN')
spark = SparkSession.builder.getOrCreate()
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
df = spark.read.format('com.databricks.spark.csv').options(header='true', inferschema='true').load('file:///home/allen/myspark/stock_small.csv')
df1 = pd.DataFrame(df.toPandas(),columns=columns_dict.keys())
firstyear = 2009
volume = 0
totol_volume_year = pd.DataFrame(columns=['year','stock_symbol','stock_volume'])
rownum = 0
for row in df1.itertuples():
    name = row[2]
    if row[1].year == firstyear and row[2] == name:        
        volume += int(row[3])
    else:
        totol_volume_year.loc[rownum,'year'] = firstyear
        totol_volume_year.loc[rownum,'stock_symbol'] = name
        totol_volume_year.loc[rownum,'stock_volume'] = volume   
        rownum+=1     
        firstyear = row[1].year
        volume = int(row[3])
totol_volume_year.loc[rownum,'year'] = firstyear
totol_volume_year.loc[rownum,'stock_symbol'] = name
totol_volume_year.loc[rownum,'stock_volume'] = volume      
groups = totol_volume_year.groupby('year')
n = np.arange(2000,2010,1)
for i in n:
    if i == 2000:
        duty11_re = groups.get_group(i).sort_values(by="stock_volume",ascending=False)
    else:
        duty11_re = duty11_re.append(groups.get_group(i).sort_values(by="stock_volume",ascending=False))

columns_json_str = '{"exchange":"exchange","stock_symbol":"stock_symbol","date":"date","stock_price_close":"stock_price_close","stock_price_open":"stock_price_open"}'
columns_dict = json.loads(columns_json_str)
df2 = pd.DataFrame(df.toPandas(),columns=columns_dict.keys())
for row in df2.itertuples():
    df2.loc[row[0],'difference'] = float(row[4]) - float(row[5])
duty12_re =df2.sort_values(by=['difference'],ascending=False)
print('------------the result of question1,duty 1--------------------')
print(duty11_re)

print('------------the result of question2,duty 1--------------------')
print(duty12_re[0:10])